{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s_qNSzzyaCbD"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jmjh290raIky"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J0Qjg6vuaHNt"
      },
      "source": [
        "# Neural machine translation with attention"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AOpGoE2T-YXS"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/text/tutorials/nmt_with_attention\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003e\n",
        "    View on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/text/blob/master/docs/tutorials/nmt_with_attention.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e\n",
        "    Run in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/text/blob/master/docs/tutorials/nmt_with_attention.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e\n",
        "    View source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/text/docs/tutorials/nmt_with_attention.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CiwtNgENbx2g"
      },
      "source": [
        "This notebook trains a sequence to sequence (seq2seq) model for Spanish to English translation based on [Effective Approaches to Attention-based Neural Machine Translation](https://arxiv.org/abs/1508.04025v5). This is an advanced example that assumes some knowledge of:\n",
        "\n",
        "* Sequence to sequence models\n",
        "* TensorFlow fundamentals below the keras layer: \n",
        "  * Working with tensors directly \n",
        "  * Writing custom `keras.Model`s and `keras.layers`\n",
        "\n",
        "While this architecture is somewhat outdated it is still a very useful project to work through to get a deeper understanding of attention mechanisms (before going on to [Transformers](transformer.ipynb)).\n",
        "\n",
        "After training the model in this notebook, you will be able to input a Spanish sentence, such as \"*¿todavia estan en casa?*\", and return the English translation: \"*are you still at home?*\"\n",
        "\n",
        "The resulting model is exportable as a `tf.saved_model`, so it can be used in other TensorFlow environments.\n",
        "\n",
        "The translation quality is reasonable for a toy example, but the generated attention plot is perhaps more interesting. This shows which parts of the input sentence has the model's attention while translating:\n",
        "\n",
        "\u003cimg src=\"https://tensorflow.org/images/spanish-english.png\" alt=\"spanish-english attention plot\"\u003e\n",
        "\n",
        "Note: This example takes approximately 10 minutes to run on a single P100 GPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yAmSR1FaqKrl"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DGFTkuRvzWqc"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tnxXKDjq3jEL"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "import typing\n",
        "from typing import Any, Tuple\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras.layers.experimental import preprocessing\n",
        "\n",
        "import tensorflow_text as tf_text\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.ticker as ticker"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vs8zge-RUdC2"
      },
      "source": [
        "This tutorial builds a few layers from scratch, use this variable if you want to switch between the custom and builtin implementations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KPJ9J7iPUchc"
      },
      "outputs": [],
      "source": [
        "use_builtins = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l_yq8kvIqoqQ"
      },
      "source": [
        "This tutorial uses a lot of low level API's where it's easy to get shapes wrong. This class is used to check shapes throughout the tutorial. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KqFqKi4fqN9X"
      },
      "outputs": [],
      "source": [
        "#@title Shape checker\n",
        "class ShapeChecker():\n",
        "  def __init__(self):\n",
        "    # Keep a cache of every axis-name seen\n",
        "    self.shapes = {}\n",
        "\n",
        "  def __call__(self, tensor, names, broadcast=False):\n",
        "    if not tf.executing_eagerly():\n",
        "      return\n",
        "\n",
        "    if isinstance(names, str):\n",
        "      names = (names,)\n",
        "\n",
        "    shape = tf.shape(tensor)\n",
        "    rank = tf.rank(tensor)\n",
        "\n",
        "    if rank != len(names):\n",
        "      raise ValueError(f'Rank mismatch:\\n'\n",
        "                       f'    found {rank}: {shape.numpy()}\\n'\n",
        "                       f'    expected {len(names)}: {names}\\n')\n",
        "\n",
        "    for i, name in enumerate(names):\n",
        "      if isinstance(name, int):\n",
        "        old_dim = name\n",
        "      else:\n",
        "        old_dim = self.shapes.get(name, None)\n",
        "      new_dim = shape[i]\n",
        "\n",
        "      if (broadcast and new_dim == 1):\n",
        "        continue\n",
        "\n",
        "      if old_dim is None:\n",
        "        # If the axis name is new, add its length to the cache.\n",
        "        self.shapes[name] = new_dim\n",
        "        continue\n",
        "\n",
        "      if new_dim != old_dim:\n",
        "        raise ValueError(f\"Shape mismatch for dimension: '{name}'\\n\"\n",
        "                         f\"    found: {new_dim}\\n\"\n",
        "                         f\"    expected: {old_dim}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gjUROhJfH3ML"
      },
      "source": [
        "## The data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puE_K74DIE9W"
      },
      "source": [
        "We'll use a language dataset provided by http://www.manythings.org/anki/. This dataset contains language translation pairs in the format:\n",
        "\n",
        "```\n",
        "May I borrow this book?\t¿Puedo tomar prestado este libro?\n",
        "```\n",
        "\n",
        "They have a variety of languages available, but we'll use the English-Spanish dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfodePkj3jEa"
      },
      "source": [
        "### Download and prepare the dataset\n",
        "\n",
        "For convenience, we've hosted a copy of this dataset on Google Cloud, but you can also download your own copy. After downloading the dataset, here are the steps we'll take to prepare the data:\n",
        "\n",
        "1. Add a *start* and *end* token to each sentence.\n",
        "2. Clean the sentences by removing special characters.\n",
        "3. Create a word index and reverse word index (dictionaries mapping from word → id and id → word).\n",
        "4. Pad each sentence to a maximum length."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kRVATYOgJs1b"
      },
      "outputs": [],
      "source": [
        "# Download the file\n",
        "import pathlib\n",
        "\n",
        "path_to_zip = tf.keras.utils.get_file(\n",
        "    'spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',\n",
        "    extract=True)\n",
        "\n",
        "path_to_file = pathlib.Path(path_to_zip).parent/'spa-eng/spa.txt'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHn4Dct23jEm"
      },
      "outputs": [],
      "source": [
        "def load_data(path):\n",
        "  text = path.read_text(encoding='utf-8')\n",
        "\n",
        "  lines = text.splitlines()\n",
        "  pairs = [line.split('\\t') for line in lines]\n",
        "\n",
        "  inp = [inp for targ, inp in pairs]\n",
        "  targ = [targ for targ, inp in pairs]\n",
        "\n",
        "  return targ, inp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cTbSbBz55QtF"
      },
      "outputs": [],
      "source": [
        "targ, inp = load_data(path_to_file)\n",
        "print(inp[-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lH_dPY8TRp3c"
      },
      "outputs": [],
      "source": [
        "print(targ[-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rgCLkfv5uO3d"
      },
      "source": [
        "### Create a tf.data dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PfVWx3WaI5Df"
      },
      "source": [
        "From these arrays of strings you can create a `tf.data.Dataset` of strings that shuffles and batches them efficiently:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TqHsArVZ3jFS"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = len(inp)\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "dataset = tf.data.Dataset.from_tensor_slices((inp, targ)).shuffle(BUFFER_SIZE)\n",
        "dataset = dataset.batch(BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qc6-NK1GtWQt"
      },
      "outputs": [],
      "source": [
        "for example_input_batch, example_target_batch in dataset.take(1):\n",
        "  print(example_input_batch[:5])\n",
        "  print()\n",
        "  print(example_target_batch[:5])\n",
        "  break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zCoxLcuN3bwv"
      },
      "source": [
        "### Text preprocessing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7kwdPcHvzz_a"
      },
      "source": [
        "One of the goals of this tutorial is to build a model that can be exported as a `tf.saved_model`. To make that exported model useful it should take `tf.string` inputs, and retrun `tf.string` outputs: All the text processing happens inside the model. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EOQ5n55X4uDB"
      },
      "source": [
        "#### Standardization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "upKhKAMK4zzI"
      },
      "source": [
        "The model is dealing with multilingual text with a limited vocabulary. So it will be important to standardize the input text.\n",
        "\n",
        "The first step is Unicode normalization to split accented characters and replace compatibility characters with their ASCII equivalents.\n",
        "\n",
        "The `tensroflow_text` package contains a unicode normalize operation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mD0e-DWGQ2Vo"
      },
      "outputs": [],
      "source": [
        "example_text = tf.constant('¿Todavía está en casa?')\n",
        "\n",
        "print(example_text.numpy())\n",
        "print(tf_text.normalize_utf8(example_text, 'NFKD').numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6hTllEjK6RSo"
      },
      "source": [
        "Unicode normalization will be the first step in the text standardization function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "chTF5N885F0P"
      },
      "outputs": [],
      "source": [
        "def tf_lower_and_split_punct(text):\n",
        "  # Split accecented characters.\n",
        "  text = tf_text.normalize_utf8(text, 'NFKD')\n",
        "  text = tf.strings.lower(text)\n",
        "  # Keep space, a to z, and select punctuation.\n",
        "  text = tf.strings.regex_replace(text, '[^ a-z.?!,¿]', '')\n",
        "  # Add spaces around punctuation.\n",
        "  text = tf.strings.regex_replace(text, '[.?!,¿]', r' \\0 ')\n",
        "  # Strip whitespace.\n",
        "  text = tf.strings.strip(text)\n",
        "\n",
        "  text = tf.strings.join(['[START]', text, '[END]'], separator=' ')\n",
        "  return text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UREvDg3sEKYa"
      },
      "outputs": [],
      "source": [
        "print(example_text.numpy().decode())\n",
        "print(tf_lower_and_split_punct(example_text).numpy().decode())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4q-sKsSI7xRZ"
      },
      "source": [
        "#### Text Vectorization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6aKn8qd37abi"
      },
      "source": [
        "This standardization function will be wrapped up in a `preprocessing.TextVectorization` layer which will handle the vocabulary extraction and conversion of input text to sequences of tokens."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eAY9k49G3jE_"
      },
      "outputs": [],
      "source": [
        "max_vocab_size = 5000\n",
        "\n",
        "input_text_processor = preprocessing.TextVectorization(\n",
        "    standardize=tf_lower_and_split_punct,\n",
        "    max_tokens=max_vocab_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7kbC6ODP8IK_"
      },
      "source": [
        "The `TextVectorization` layer and many other `experimental.preprocessing` layers have an `adapt` method. This method reads one epoch of the training data, and works a lot like `Model.fix`. This `adapt` method initializes the layer based on the data. Here it determines the vocabulary:  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bmsI1Yql8FYe"
      },
      "outputs": [],
      "source": [
        "input_text_processor.adapt(inp)\n",
        "\n",
        "# Here are the first 10 words from the vocabulary:\n",
        "input_text_processor.get_vocabulary()[:10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9kGjIFjX8_Wp"
      },
      "source": [
        "That's the Spanish `TextVectorization` layer, now build and `.adapt()` the English one:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jlC4xuZnKLBS"
      },
      "outputs": [],
      "source": [
        "output_text_processor = preprocessing.TextVectorization(\n",
        "    standardize=tf_lower_and_split_punct,\n",
        "    max_tokens=max_vocab_size)\n",
        "\n",
        "output_text_processor.adapt(targ)\n",
        "output_text_processor.get_vocabulary()[:10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BWQqlP_s9eIv"
      },
      "source": [
        "Now these layers can convert a batch of strings into a batch of token IDs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9KZxj8IrNZ9S"
      },
      "outputs": [],
      "source": [
        "example_tokens = input_text_processor(example_input_batch)\n",
        "example_tokens[:3, :10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AA9rUn9G9n78"
      },
      "source": [
        "The `get_vocabulary` method can be used to convert token IDs back to text:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "98g9rcxGQY0I"
      },
      "outputs": [],
      "source": [
        "input_vocab = np.array(input_text_processor.get_vocabulary())\n",
        "tokens = input_vocab[example_tokens[0].numpy()]\n",
        "' '.join(tokens)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ot0aCL9t-Ghi"
      },
      "source": [
        "The returned token IDs are zero-padded. This can easily be turned into a mask:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_jx4Or_eFRSz"
      },
      "outputs": [],
      "source": [
        "plt.subplot(1, 2, 1)\n",
        "plt.pcolormesh(example_tokens)\n",
        "plt.title('Token IDs')\n",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.pcolormesh(example_tokens != 0)\n",
        "plt.title('Mask')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TNfHIF71ulLu"
      },
      "source": [
        "## The encoder/decoder model\n",
        "\n",
        "The following diagram shows an overview of the model. At each time-step the decoder's output is combined with a weighted sum over the encoded input, to predict the next word. The diagram and formulas are from [Luong's paper](https://arxiv.org/abs/1508.04025v5).\n",
        "\n",
        "\u003cimg src=\"https://www.tensorflow.org/images/seq2seq/attention_mechanism.jpg\" width=\"500\" alt=\"attention mechanism\"\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gzQWx2saImMV"
      },
      "source": [
        "Before getting into it define a few constants for the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_a9uNz3-IrF-"
      },
      "outputs": [],
      "source": [
        "embedding_dim = 256\n",
        "units = 1024"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "blNgVbLSzpsr"
      },
      "source": [
        "### The encoder\n",
        "\n",
        "Start by building the encoder, the blue part of the diagram above.\n",
        "\n",
        "The encoder:\n",
        "\n",
        "1. Takes a list of token IDs (from `input_text_processor`).\n",
        "3. Looks up an embedding vector for each token (Using a `layers.Embedding`).\n",
        "4. Processes the embeddings into a new sequence (Using a `layers.GRU`).\n",
        "5. Returns:\n",
        "  * The processed sequence. This will be passed to the attention head.\n",
        "  * The internal state. This will be used to initialize the decoder\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nZ2rI24i3jFg"
      },
      "outputs": [],
      "source": [
        "class Encoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, input_vocab_size, embedding_dim, enc_units):\n",
        "    super(Encoder, self).__init__()\n",
        "    self.enc_units = enc_units\n",
        "    self.input_vocab_size = input_vocab_size\n",
        "\n",
        "    # The embedding layer converts tokens to vectors\n",
        "    self.embedding = tf.keras.layers.Embedding(self.input_vocab_size,\n",
        "                                               embedding_dim)\n",
        "\n",
        "    # The GRU RNN layer processes those vectors sequentially.\n",
        "    self.gru = tf.keras.layers.GRU(self.enc_units,\n",
        "                                   # Return the sequence and state\n",
        "                                   return_sequences=True,\n",
        "                                   return_state=True,\n",
        "                                   recurrent_initializer='glorot_uniform')\n",
        "\n",
        "  def call(self, tokens, state=None):\n",
        "    shape_checker = ShapeChecker()\n",
        "    shape_checker(tokens, ('batch', 's'))\n",
        "\n",
        "    # 2. The embedding layer looks up the embedding for each token.\n",
        "    vectors = self.embedding(tokens)\n",
        "    shape_checker(vectors, ('batch', 's', 'embed_dim'))\n",
        "\n",
        "    # 3. The GRU processes the embedding sequence.\n",
        "    #    output shape: (batch, s, enc_units)\n",
        "    #    state shape: (batch, enc_units)\n",
        "    output, state = self.gru(vectors, initial_state=state)\n",
        "    shape_checker(output, ('batch', 's', 'enc_units'))\n",
        "    shape_checker(state, ('batch', 'enc_units'))\n",
        "\n",
        "    # 4. Returns the new sequence and its state.\n",
        "    return output, state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D3SKkaQeGn-Q"
      },
      "source": [
        "Here is how it fits together so far:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "60gSVh05Jl6l"
      },
      "outputs": [],
      "source": [
        "# Convert the input text to tokens.\n",
        "example_tokens = input_text_processor(example_input_batch)\n",
        "\n",
        "# Encode the input sequence.\n",
        "encoder = Encoder(input_text_processor.vocabulary_size(),\n",
        "                  embedding_dim, units)\n",
        "example_enc_output, example_enc_state = encoder(example_tokens)\n",
        "\n",
        "print(f'Input batch, shape (batch): {example_input_batch.shape}')\n",
        "print(f'Input batch tokens, shape (batch, s): {example_tokens.shape}')\n",
        "print(f'Encoder output, shape (batch, s, units): {example_enc_output.shape}')\n",
        "print(f'Encoder state, shape (batch, units): {example_enc_state.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RIPHh4O9ixB"
      },
      "source": [
        "The encoder returns its internal state so that its state can be used to initialize the decoder.\n",
        "\n",
        "It's also common for an RNN to return its state so that it can process a sequence over multiple calls. You'll see more of that building the decoder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45xM_Gl1MgXY"
      },
      "source": [
        "### The attention head\n",
        "\n",
        "The decoder uses attention to selectively focus on parts of the input sequence.\n",
        "The attention takes a sequence of vectors as input for each example and returns an \"attention\" vector for each example. This attention layer is similar to a `layers.GlobalAveragePoling1D` but the attention layer performs a _weighted_ average.\n",
        "\n",
        "Let's look at how this works:\n",
        "\n",
        "\u003cimg src=\"images/attention_equation_1.jpg\" alt=\"attention equation 1\" width=\"800\"\u003e\n",
        "\n",
        "\u003cimg src=\"images/attention_equation_2.jpg\" alt=\"attention equation 2\" width=\"800\"\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NX2JsKzzzgZ5"
      },
      "source": [
        "Where:\n",
        "\n",
        "* $s$ is the encoder index.\n",
        "* $t$ is the decoder index.\n",
        "* $\\alpha_{ts}$ is the attention weights.\n",
        "* $h_s$ is the sequence of encoder outputs being attended to (the attention \"key\" and \"value\" in transformer terminology).\n",
        "* $h_t$ is the the decoder state attending to the sequence (the attention \"query\" in transformer terminology).\n",
        "* $c_t$ is the resulting context vector.\n",
        "* $a_t$ is the final output combining the \"context\" and \"query\".\n",
        "\n",
        "The equations:\n",
        "\n",
        "1. Calculates the attention weights, $\\alpha_{ts}$, as a softmax across the encoder's output sequence.\n",
        "2. Calculates the context vector as the weighted sum of the encoder outputs.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fNA5GeHHPsGL"
      },
      "source": [
        "Last is the $score$ function. Its job is to calculate a scalar logit-score for each key-query pair. There are two common approaches:\n",
        "\n",
        "\u003cimg src=\"images/attention_equation_4.jpg\" alt=\"attention equation 4\" width=\"800\"\u003e\n",
        "\n",
        "This tutorial uses [Bahdanau's additive attention](https://arxiv.org/pdf/1409.0473.pdf). TensorFlow includes implementations of both as `layers.Attention` and\n",
        "`layers.AdditiveAttention`. The class below handles the weight matrices in a pair of `layers.Dense` layers, and calls the builtin implementation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "momiE59lXo6U"
      },
      "outputs": [],
      "source": [
        "class BahdanauAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, units):\n",
        "    super().__init__()\n",
        "    # For Eqn. (4), the  Bahdanau attention\n",
        "    self.W1 = tf.keras.layers.Dense(units, use_bias=False)\n",
        "    self.W2 = tf.keras.layers.Dense(units, use_bias=False)\n",
        "\n",
        "    self.attention = tf.keras.layers.AdditiveAttention()\n",
        "\n",
        "  def call(self, query, value, mask):\n",
        "    shape_checker = ShapeChecker()\n",
        "    shape_checker(query, ('batch', 't', 'query_units'))\n",
        "    shape_checker(value, ('batch', 's', 'value_units'))\n",
        "    shape_checker(mask, ('batch', 's'))\n",
        "\n",
        "    # From Eqn. (4), `W1@ht`.\n",
        "    w1_query = self.W1(query)\n",
        "    shape_checker(w1_query, ('batch', 't', 'attn_units'))\n",
        "\n",
        "    # From Eqn. (4), `W2@hs`.\n",
        "    w2_key = self.W2(value)\n",
        "    shape_checker(w2_key, ('batch', 's', 'attn_units'))\n",
        "\n",
        "    query_mask = tf.ones(tf.shape(query)[:-1], dtype=bool)\n",
        "    value_mask = mask\n",
        "\n",
        "    context_vector, attention_weights = self.attention(\n",
        "        inputs = [w1_query, value, w2_key],\n",
        "        mask=[query_mask, value_mask],\n",
        "        return_attention_scores = True,\n",
        "    )\n",
        "    shape_checker(context_vector, ('batch', 't', 'value_units'))\n",
        "    shape_checker(attention_weights, ('batch', 't', 's'))\n",
        "\n",
        "    return context_vector, attention_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cf13LubPGjDO"
      },
      "source": [
        "### Test the Attention layer\n",
        "\n",
        "Create a `BahdanauAttention` layer:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t4QMlOp8Gidh"
      },
      "outputs": [],
      "source": [
        "attention_layer = BahdanauAttention(units)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "snA1uL9AI-JE"
      },
      "source": [
        "This layer takes 3 inputs:\n",
        "\n",
        "* The `query`: This will be generated by the decoder, later.\n",
        "* The `value`: This Will be the output of the encoder.\n",
        "* The `mask`: To exclude the padding, `example_tokens != 0`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DYSHqmORgVFo"
      },
      "outputs": [],
      "source": [
        "(example_tokens != 0).shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g2bmvT25pXnr"
      },
      "source": [
        "The vectorized implementation of the attention layer lets you pass a batch of sequences of query vectors and a batch of sequence of value vectors. The result is:\n",
        "\n",
        "1. A batch of sequences of result vectors the size of the queries.\n",
        "2. A batch attention maps, with size `(query_length, value_length)`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7y7hjPkNMmHh"
      },
      "outputs": [],
      "source": [
        "# Later, the decoder will generate this attention query\n",
        "example_attention_query = tf.random.normal(shape=[len(example_tokens), 2, 10])\n",
        "\n",
        "# Attend to the encoded tokens\n",
        "\n",
        "context_vector, attention_weights = attention_layer(\n",
        "    query=example_attention_query,\n",
        "    value=example_enc_output,\n",
        "    mask=(example_tokens != 0))\n",
        "\n",
        "print(f'Attention result shape: (batch_size, query_seq_length, units):           {context_vector.shape}')\n",
        "print(f'Attention weights shape: (batch_size, query_seq_length, value_seq_length): {attention_weights.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AagyXMH-Jhqt"
      },
      "source": [
        "The attention weights should sum to `1.0` for each sequence.\n",
        "\n",
        "Here are the attention weights across the sequences at `t=0`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rqr8XGsAJlf6"
      },
      "outputs": [],
      "source": [
        "plt.subplot(1, 2, 1)\n",
        "plt.pcolormesh(attention_weights[:, 0, :])\n",
        "plt.title('Attention weights')\n",
        "\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.pcolormesh(example_tokens != 0)\n",
        "plt.title('Mask')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Eil-C_NN1rp"
      },
      "source": [
        "Because of the small-random initialization the attention weights are all close to `1/(sequence_length)`. If you zoom in on the weights for a single sequence, you can see that there is some _small_ variation that the model can learn to expand, and exploit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZuzrCdmYlTcJ"
      },
      "outputs": [],
      "source": [
        "attention_weights.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qIMwC-f-ZC8N"
      },
      "outputs": [],
      "source": [
        "attention_slice = attention_weights[0, 0].numpy()\n",
        "attention_slice = attention_slice[attention_slice != 0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ysWDPO6hOS8X"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "plt.suptitle('Attention weights for one sequence')\n",
        "\n",
        "plt.figure(figsize=(12, 6))\n",
        "a1 = plt.subplot(1, 2, 1)\n",
        "plt.bar(range(len(attention_slice)), attention_slice)\n",
        "# freeze the xlim\n",
        "plt.xlim(plt.xlim())\n",
        "plt.xlabel('Attention weights')\n",
        "\n",
        "a2 = plt.subplot(1, 2, 2)\n",
        "plt.bar(range(len(attention_slice)), attention_slice)\n",
        "plt.xlabel('Attention weights, zoomed')\n",
        "\n",
        "# zoom in\n",
        "top = max(a1.get_ylim())\n",
        "zoom = 0.85*top\n",
        "a2.set_ylim([0.90*top, top])\n",
        "a1.plot(a1.get_xlim(), [zoom, zoom], color='k')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQ638eHN4iCK"
      },
      "source": [
        "### The decoder\n",
        "\n",
        "The decoder's job is to generate predictions for the next output token.\n",
        "\n",
        "1. The decoder receives the complete encoder output.\n",
        "2. It uses an RNN to keep track of what it has generated so far.\n",
        "3. It uses its RNN output as the query to the attention over the encoder's output, producing the context vector.\n",
        "4. It combines the RNN output and the context vector using Equation 3 (below) to generate the \"attention vector\".\n",
        "5. It generates logit predictions for the next token based on the \"attention vector\".\n",
        "\n",
        "\u003cimg src=\"images/attention_equation_3.jpg\" alt=\"attention equation 3\" width=\"800\"\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pZsQJMqNmg_L"
      },
      "source": [
        "Here is the `Decoder` class and its initializer. The initializer creates all the necessary layers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "erYvHIgAl8kh"
      },
      "outputs": [],
      "source": [
        "class Decoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, output_vocab_size, embedding_dim, dec_units):\n",
        "    super(Decoder, self).__init__()\n",
        "    self.dec_units = dec_units\n",
        "    self.output_vocab_size = output_vocab_size\n",
        "    self.embedding_dim = embedding_dim\n",
        "\n",
        "    # For Step 1. The embedding layer convets token IDs to vectors\n",
        "    self.embedding = tf.keras.layers.Embedding(self.output_vocab_size,\n",
        "                                               embedding_dim)\n",
        "\n",
        "    # For Step 2. The RNN keeps track of what's been generated so far.\n",
        "    self.gru = tf.keras.layers.GRU(self.dec_units,\n",
        "                                   return_sequences=True,\n",
        "                                   return_state=True,\n",
        "                                   recurrent_initializer='glorot_uniform')\n",
        "\n",
        "    # For step 3. The RNN output will be the query for the attention layer.\n",
        "    self.attention = BahdanauAttention(self.dec_units)\n",
        "\n",
        "    # For step 4. Eqn. (3): converting `ct` to `at`\n",
        "    self.Wc = tf.keras.layers.Dense(dec_units, activation=tf.math.tanh,\n",
        "                                    use_bias=False)\n",
        "\n",
        "    # For step 5. This fully connected layer produces the logits for each\n",
        "    # output token.\n",
        "    self.fc = tf.keras.layers.Dense(self.output_vocab_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eUTfYHmfmwKH"
      },
      "source": [
        "The `call` method for this layer takes and returns multiple tensors. Organize those into simple container classes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7WfSIb2sArRT"
      },
      "outputs": [],
      "source": [
        "class DecoderInput(typing.NamedTuple):\n",
        "  new_tokens: Any\n",
        "  enc_output: Any\n",
        "  mask: Any\n",
        "\n",
        "class DecoderOutput(typing.NamedTuple):\n",
        "  logits: Any\n",
        "  attention_weights: Any"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NChkl2KrnV2y"
      },
      "source": [
        "Here is the implementation of the `call` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJOi5btHAPNK"
      },
      "outputs": [],
      "source": [
        "def call(self,\n",
        "         inputs: DecoderInput,\n",
        "         state=None) -\u003e Tuple[DecoderOutput, tf.Tensor]:\n",
        "  shape_checker = ShapeChecker()\n",
        "  shape_checker(inputs.new_tokens, ('batch', 't'))\n",
        "  shape_checker(inputs.enc_output, ('batch', 's', 'enc_units'))\n",
        "  shape_checker(inputs.mask, ('batch', 's'))\n",
        "\n",
        "  if state is not None:\n",
        "    shape_checker(state, ('batch', 'dec_units'))\n",
        "\n",
        "  # Step 1. Lookup the embeddings\n",
        "  vectors = self.embedding(inputs.new_tokens)\n",
        "  shape_checker(vectors, ('batch', 't', 'embedding_dim'))\n",
        "\n",
        "  # Step 2. Process one step with the RNN\n",
        "  rnn_output, state = self.gru(vectors, initial_state=state)\n",
        "\n",
        "  shape_checker(rnn_output, ('batch', 't', 'dec_units'))\n",
        "  shape_checker(state, ('batch', 'dec_units'))\n",
        "\n",
        "  # Step 3. Use the RNN output as the query for the attention over the\n",
        "  # encoder output.\n",
        "  context_vector, attention_weights = self.attention(\n",
        "      query=rnn_output, value=inputs.enc_output, mask=inputs.mask)\n",
        "  shape_checker(context_vector, ('batch', 't', 'dec_units'))\n",
        "  shape_checker(attention_weights, ('batch', 't', 's'))\n",
        "\n",
        "  # Step 4. Eqn. (3): Join the context_vector and rnn_output\n",
        "  #     [ct; ht] shape: (batch t, value_units + query_units)\n",
        "  context_and_rnn_output = tf.concat([context_vector, rnn_output], axis=-1)\n",
        "\n",
        "  # Step 4. Eqn. (3): `at = tanh(Wc@[ct; ht])`\n",
        "  attention_vector = self.Wc(context_and_rnn_output)\n",
        "  shape_checker(attention_vector, ('batch', 't', 'dec_units'))\n",
        "\n",
        "  # Step 5. Generate logit predictions:\n",
        "  logits = self.fc(attention_vector)\n",
        "  shape_checker(logits, ('batch', 't', 'output_vocab_size'))\n",
        "\n",
        "  return DecoderOutput(logits, attention_weights), state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ay_mTMPfnb2a"
      },
      "outputs": [],
      "source": [
        "Decoder.call = call"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "arTOBklcFTiC"
      },
      "source": [
        "The **encoder** processes its full input sequence with a single call to its RNN. This implementation of the **decoder** _can_ do that as well for efficient training. But this tutorial will run the decoder in a loop for a few reasons:\n",
        "\n",
        "* Flexibility: Writing the loop gives you direct control over the training procedure.\n",
        "* Clarity: It's possible to do masking tricks and use `layers.RNN`, or `tfa.seq2seq` APIs to pack this all into a single call. But writing it out as a loop may be clearer. \n",
        "  * Loop free training is demonstrated in the [Text generation](text_generation.ipynb) tutiorial.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E1-mLAcUEXpK"
      },
      "source": [
        "Now try using this decoder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ZUMbYXIEVeA"
      },
      "outputs": [],
      "source": [
        "decoder = Decoder(output_text_processor.vocabulary_size(),\n",
        "                  embedding_dim, units)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UPnaw583CpnY"
      },
      "source": [
        "The decoder takes 4 inputs.\n",
        "\n",
        "* `new_tokens` -  The last token generated. Initialize the decoder with the `\"[START]\"` token.\n",
        "* `enc_output` - Generated by the `Encoder`.\n",
        "* `mask` - A boolean tensor indicating where `tokens != 0`\n",
        "* `state` - The previous `state` output from the decoder (the internal state\n",
        "  of the decoder's RNN). Pass `None` to zero-initialize it. The original\n",
        "  paper initializes it from the encoder's final RNN state. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4u6eJBU4GL40"
      },
      "outputs": [],
      "source": [
        "# Convert the target sequence, and collect the \"[START]\" tokens\n",
        "example_output_tokens = output_text_processor(example_target_batch)\n",
        "\n",
        "start_index = output_text_processor.get_vocabulary().index('[START]')\n",
        "first_token = tf.constant([[start_index]] * example_output_tokens.shape[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E5hqvbR5FUCD"
      },
      "outputs": [],
      "source": [
        "# Run the decoder\n",
        "dec_result, dec_state = decoder(\n",
        "    inputs = DecoderInput(new_tokens=first_token,\n",
        "                          enc_output=example_enc_output,\n",
        "                          mask=(example_tokens != 0)),\n",
        "    state = example_enc_state\n",
        ")\n",
        "\n",
        "print(f'logits shape: (batch_size, t, output_vocab_size) {dec_result.logits.shape}')\n",
        "print(f'state shape: (batch_size, dec_units) {dec_state.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vEZvXZRVPHd6"
      },
      "source": [
        "Sample a token according to the logits:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5UY8wko3jFp"
      },
      "outputs": [],
      "source": [
        "sampled_token = tf.random.categorical(dec_result.logits[:, 0, :], num_samples=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-xTpX44VkzrY"
      },
      "source": [
        "Decode the token as the first word of the output:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lKXTLYu4IV7I"
      },
      "outputs": [],
      "source": [
        "vocab = np.array(output_text_processor.get_vocabulary())\n",
        "first_word = vocab[sampled_token.numpy()]\n",
        "first_word[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LUQV6AXoQR7z"
      },
      "source": [
        "Now use the decoder to generate a second set of logits.\n",
        "\n",
        "- Pass the same `enc_output` and `mask`, these haven't changed.\n",
        "- Pass the sampled token as `new_tokens`.\n",
        "- Pass the `decoder_state` the decoder returned last time, so the RNN continues with a memory of where it left off last time.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pX1VF9XDJTOM"
      },
      "outputs": [],
      "source": [
        "dec_result, dec_state = decoder(\n",
        "    DecoderInput(sampled_token,\n",
        "                 example_enc_output,\n",
        "                 mask=(example_tokens != 0)),\n",
        "    state=dec_state)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H1rs0XL7Y2aS"
      },
      "outputs": [],
      "source": [
        "sampled_token = tf.random.categorical(dec_result.logits[:, 0, :], num_samples=1)\n",
        "first_word = vocab[sampled_token.numpy()]\n",
        "first_word[:5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B6xyru86m914"
      },
      "source": [
        "## Training\n",
        "\n",
        "Now that you have all the model components, it's time to start training the model. You'll need:\n",
        "\n",
        "- A loss function and optimizer to perform the optimization.\n",
        "- A training step function defining how to update the model for each input/target batch.\n",
        "- A training loop to drive the training and save checkpoints."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ch_71VbIRfK"
      },
      "source": [
        "### Define the loss function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WmTHr5iV3jFr"
      },
      "outputs": [],
      "source": [
        "class MaskedLoss(tf.keras.losses.Loss):\n",
        "  def __init__(self):\n",
        "    self.name = 'masked_loss'\n",
        "    self.loss = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "        from_logits=True, reduction='none')\n",
        "\n",
        "  def __call__(self, y_true, y_pred):\n",
        "    shape_checker = ShapeChecker()\n",
        "    shape_checker(y_true, ('batch', 't'))\n",
        "    shape_checker(y_pred, ('batch', 't', 'logits'))\n",
        "\n",
        "    # Calculate the loss for each item in the batch.\n",
        "    loss = self.loss(y_true, y_pred)\n",
        "    shape_checker(loss, ('batch', 't'))\n",
        "\n",
        "    # Mask off the losses on padding.\n",
        "    mask = tf.cast(y_true != 0, tf.float32)\n",
        "    shape_checker(mask, ('batch', 't'))\n",
        "    loss *= mask\n",
        "\n",
        "    # Return the total.\n",
        "    return tf.reduce_sum(loss)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M5AgEBh2S404"
      },
      "source": [
        "### Implement the training step"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r_G20Te1XSmJ"
      },
      "source": [
        "Start with a model class, the training process will be implemented as the `train_step` method on this model. See [Customizing fit](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit) for details.\n",
        "\n",
        "Here the `train_step` method is a wrapper around the `_train_step` implementation which will come later. This wrapper includes a switch to turn on and off `tf.function` compilation, to make debugging easier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WWIyuy71TkJT"
      },
      "outputs": [],
      "source": [
        "class TrainTranslator(tf.keras.Model):\n",
        "  def __init__(self, embedding_dim, units,\n",
        "               input_text_processor,\n",
        "               output_text_processor, \n",
        "               use_tf_function=True):\n",
        "    super().__init__()\n",
        "    # Build the encoder and decoder\n",
        "    encoder = Encoder(input_text_processor.vocabulary_size(),\n",
        "                      embedding_dim, units)\n",
        "    decoder = Decoder(output_text_processor.vocabulary_size(),\n",
        "                      embedding_dim, units)\n",
        "\n",
        "    self.encoder = encoder\n",
        "    self.decoder = decoder\n",
        "    self.input_text_processor = input_text_processor\n",
        "    self.output_text_processor = output_text_processor\n",
        "    self.use_tf_function = use_tf_function\n",
        "    self.shape_checker = ShapeChecker()\n",
        "\n",
        "  def train_step(self, inputs):\n",
        "    self.shape_checker = ShapeChecker()\n",
        "    if self.use_tf_function:\n",
        "      return self._tf_train_step(inputs)\n",
        "    else:\n",
        "      return self._train_step(inputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-i0i1x6jwsLm"
      },
      "source": [
        "Overall the implementation for the `Model.train_step` method is as follows:\n",
        "\n",
        "1. Receive a batch of `input_text, target_text` from the `tf.data.Dataset`.\n",
        "2. Convert those raw text inputs to token-embeddings and masks. \n",
        "3. Run the encoder on the `input_tokens` to get the `encoder_output` and `encoder_state`.\n",
        "4. Initialize the decoder state and loss. \n",
        "5. Loop over the `target_tokens`:\n",
        "   1. Run the decoder one step at a time.\n",
        "   2. Calculate the loss for each step.\n",
        "   3. Accumulate the average loss.\n",
        "6. Calculate the gradient of the loss and use the optimizer to apply updates to the model's `trainable_variables`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ngBjFw4BU5G7"
      },
      "source": [
        "The `_preprocess` method, added below, implements steps #1 and #2: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZlYE68wzXoA8"
      },
      "outputs": [],
      "source": [
        "def _preprocess(self, input_text, target_text):\n",
        "  self.shape_checker(input_text, ('batch',))\n",
        "  self.shape_checker(target_text, ('batch',))\n",
        "\n",
        "  # Convert the text to token IDs\n",
        "  input_tokens = self.input_text_processor(input_text)\n",
        "  target_tokens = self.output_text_processor(target_text)\n",
        "  self.shape_checker(input_tokens, ('batch', 's'))\n",
        "  self.shape_checker(target_tokens, ('batch', 't'))\n",
        "\n",
        "  # Convert IDs to masks.\n",
        "  input_mask = input_tokens != 0\n",
        "  self.shape_checker(input_mask, ('batch', 's'))\n",
        "\n",
        "  target_mask = target_tokens != 0\n",
        "  self.shape_checker(target_mask, ('batch', 't'))\n",
        "\n",
        "  return input_tokens, input_mask, target_tokens, target_mask"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lHy6hzStrgjQ"
      },
      "outputs": [],
      "source": [
        "TrainTranslator._preprocess = _preprocess"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d3kvbcArc2y-"
      },
      "source": [
        "The `_train_step` method, added below, handles the remaining steps except for actually running the decoder: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qs_gsISsYPpY"
      },
      "outputs": [],
      "source": [
        "def _train_step(self, inputs):\n",
        "  input_text, target_text = inputs  \n",
        "\n",
        "  (input_tokens, input_mask,\n",
        "   target_tokens, target_mask) = self._preprocess(input_text, target_text)\n",
        "\n",
        "  max_target_length = tf.shape(target_tokens)[1]\n",
        "\n",
        "  with tf.GradientTape() as tape:\n",
        "    # Encode the input\n",
        "    enc_output, enc_state = self.encoder(input_tokens)\n",
        "    self.shape_checker(enc_output, ('batch', 's', 'enc_units'))\n",
        "    self.shape_checker(enc_state, ('batch', 'enc_units'))\n",
        "\n",
        "    # Initialize the decoder's state to the encoder's final state.\n",
        "    # This only works if the encoder and decoder have the same number of\n",
        "    # units.\n",
        "    dec_state = enc_state\n",
        "    loss = tf.constant(0.0)\n",
        "\n",
        "    for t in tf.range(max_target_length-1):\n",
        "      # Pass in two tokens from the target sequence:\n",
        "      # 1. The current input to the decoder.\n",
        "      # 2. The target for the decoder's next prediction.\n",
        "      new_tokens = target_tokens[:, t:t+2]\n",
        "      step_loss, dec_state = self._loop_step(new_tokens, input_mask,\n",
        "                                             enc_output, dec_state)\n",
        "      loss = loss + step_loss\n",
        "\n",
        "    # Average the loss over all non padding tokens.\n",
        "    average_loss = loss / tf.reduce_sum(tf.cast(target_mask, tf.float32))\n",
        "\n",
        "  # Apply an optimization step\n",
        "  variables = self.trainable_variables \n",
        "  gradients = tape.gradient(average_loss, variables)\n",
        "  self.optimizer.apply_gradients(zip(gradients, variables))\n",
        "\n",
        "  # Return a dict mapping metric names to current value\n",
        "  return {'batch_loss': average_loss}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KGwWHIxLrjGR"
      },
      "outputs": [],
      "source": [
        "TrainTranslator._train_step = _train_step"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F7g40o-mXyt5"
      },
      "source": [
        "The `_loop_step` method, added below, executes the decoder and calculates the incremental loss and new decoder state (`dec_state`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9VrzgwztXzYJ"
      },
      "outputs": [],
      "source": [
        "def _loop_step(self, new_tokens, input_mask, enc_output, dec_state):\n",
        "  input_token, target_token = new_tokens[:, 0:1], new_tokens[:, 1:2]\n",
        "\n",
        "  # Run the decoder one step.\n",
        "  decoder_input = DecoderInput(new_tokens=input_token,\n",
        "                               enc_output=enc_output,\n",
        "                               mask=input_mask)\n",
        "\n",
        "  dec_result, dec_state = self.decoder(decoder_input, state=dec_state)\n",
        "  self.shape_checker(dec_result.logits, ('batch', 't1', 'logits'))\n",
        "  self.shape_checker(dec_result.attention_weights, ('batch', 't1', 's'))\n",
        "  self.shape_checker(dec_state, ('batch', 'dec_units'))\n",
        "\n",
        "  # `self.loss` returns the total for non-padded tokens\n",
        "  y = target_token\n",
        "  y_pred = dec_result.logits\n",
        "  step_loss = self.loss(y, y_pred)\n",
        "\n",
        "  return step_loss, dec_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xj3I7VULrk1R"
      },
      "outputs": [],
      "source": [
        "TrainTranslator._loop_step = _loop_step"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WACCHvKWBQ9C"
      },
      "source": [
        "### Test the training step\n",
        "\n",
        "Build a `TrainTranslator`, and configure it for training using the `Model.compile` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OA6bCske8TXm"
      },
      "outputs": [],
      "source": [
        "translator = TrainTranslator(\n",
        "    embedding_dim, units,\n",
        "    input_text_processor=input_text_processor,\n",
        "    output_text_processor=output_text_processor,\n",
        "    use_tf_function=False)\n",
        "\n",
        "# Configure the loss and optimizer\n",
        "translator.compile(\n",
        "    optimizer=tf.optimizers.Adam(),\n",
        "    loss=MaskedLoss(),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6y5OnZDsB3sB"
      },
      "source": [
        "Test out the `train_step`. For a text model like this the loss should start near:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zHe-OudqCFGK"
      },
      "outputs": [],
      "source": [
        "np.log(output_text_processor.vocabulary_size())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VwMU9cFEfjha"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "for n in range(10):\n",
        "  print(translator.train_step([example_input_batch, example_target_batch]))\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-xqtsMbCUp2"
      },
      "source": [
        "While it's easier to debug without a `tf.function` it does give a performance boost. So now that the `_train_step` method is working, try the `tf.function`-wrapped `_tf_train_step`, to maximize performance while training:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UFUsTKQx0jaH"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=[[tf.TensorSpec(dtype=tf.string, shape=[None]),\n",
        "                               tf.TensorSpec(dtype=tf.string, shape=[None])]])\n",
        "def _tf_train_step(self, inputs):\n",
        "  return self._train_step(inputs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2-bgU59jrztQ"
      },
      "outputs": [],
      "source": [
        "TrainTranslator._tf_train_step = _tf_train_step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KC8bRv_Gr3H9"
      },
      "outputs": [],
      "source": [
        "translator.use_tf_function = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EKMYNF_sIFb9"
      },
      "source": [
        "The first call will be slow, because it traces the function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pLQZsX2dp1QK"
      },
      "outputs": [],
      "source": [
        "translator.train_step([example_input_batch, example_target_batch])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W3t2Hg7UISYi"
      },
      "source": [
        "But after that it's usually 2-3x faster than the eager `train_step` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UzXXMwjXCqqh"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "for n in range(10):\n",
        "  print(translator.train_step([example_input_batch, example_target_batch]))\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OIvigTqaEcu1"
      },
      "source": [
        "A good test of a new model is to see that it can overfit a single batch of input. Try it, the loss should quickly go to zero:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U-dIWMIBqK7b"
      },
      "outputs": [],
      "source": [
        "losses = []\n",
        "for n in range(100):\n",
        "  print('.', end='')\n",
        "  logs = translator.train_step([example_input_batch, example_target_batch])\n",
        "  losses.append(logs['batch_loss'].numpy())\n",
        "\n",
        "print()\n",
        "plt.plot(losses)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aI02XFjoEt1k"
      },
      "source": [
        "Now that you're confident that the training step is working, build a fresh copy of the model to train from scratch:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Emgfgh4tAmJt"
      },
      "outputs": [],
      "source": [
        "train_translator = TrainTranslator(\n",
        "    embedding_dim, units,\n",
        "    input_text_processor=input_text_processor,\n",
        "    output_text_processor=output_text_processor)\n",
        "\n",
        "# Configure the loss and optimizer\n",
        "train_translator.compile(\n",
        "    optimizer=tf.optimizers.Adam(),\n",
        "    loss=MaskedLoss(),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpObfY22IddU"
      },
      "source": [
        "### Train the model\n",
        "\n",
        "While there's nothing wrong with writing your own custom training loop, implementing the `Model.train_step` method, as in the previous section, allows you to run `Model.fit` and avoid rewriting all that boiler-plate code. \n",
        "\n",
        "This tutorial only trains for a couple of epochs, so use a `callbacks.Callback` to collect the history of batch losses, for plotting:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J7m4mtnj80sq"
      },
      "outputs": [],
      "source": [
        "class BatchLogs(tf.keras.callbacks.Callback):\n",
        "  def __init__(self, key):\n",
        "    self.key = key\n",
        "    self.logs = []\n",
        "\n",
        "  def on_train_batch_end(self, n, logs):\n",
        "    self.logs.append(logs[self.key])\n",
        "\n",
        "batch_loss = BatchLogs('batch_loss')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BQd_esVVoSf3"
      },
      "outputs": [],
      "source": [
        "train_translator.fit(dataset, epochs=3,\n",
        "                     callbacks=[batch_loss])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "38rLdlmtQHCm"
      },
      "outputs": [],
      "source": [
        "plt.plot(batch_loss.logs)\n",
        "plt.ylim([0, 3])\n",
        "plt.xlabel('Batch #')\n",
        "plt.ylabel('CE/token')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0S_O_RzHmfe"
      },
      "source": [
        "The visible jumps in the plot are at the epoch boundaries."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mU3Ce8M6I3rz"
      },
      "source": [
        "## Translate\n",
        "\n",
        "Now that the model is trained, implement a function to execute the full `text =\u003e text` translation.\n",
        "\n",
        "For this the model needs to invert the `text =\u003e token IDs` mapping provided by the `output_text_processor`. It also needs to know the IDs for special tokens. This is all implemented in the constructor for the new class. The implementation of the actual translate method will follow.\n",
        "\n",
        "Overall this is similar to the training loop, except that the input to the decoder at each time step is a sample from the decoder's last prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PO-CLL1LVBbM"
      },
      "outputs": [],
      "source": [
        "class Translator(tf.Module):\n",
        "\n",
        "  def __init__(self, encoder, decoder, input_text_processor,\n",
        "               output_text_processor):\n",
        "    self.encoder = encoder\n",
        "    self.decoder = decoder\n",
        "    self.input_text_processor = input_text_processor\n",
        "    self.output_text_processor = output_text_processor\n",
        "\n",
        "    self.output_token_string_from_index = (\n",
        "        tf.keras.layers.experimental.preprocessing.StringLookup(\n",
        "            vocabulary=output_text_processor.get_vocabulary(),\n",
        "            mask_token='',\n",
        "            invert=True))\n",
        "\n",
        "    # The output should never generate padding, unknown, or start.\n",
        "    index_from_string = tf.keras.layers.experimental.preprocessing.StringLookup(\n",
        "        vocabulary=output_text_processor.get_vocabulary(), mask_token='')\n",
        "    token_mask_ids = index_from_string(['', '[UNK]', '[START]']).numpy()\n",
        "\n",
        "    token_mask = np.zeros([index_from_string.vocabulary_size()], dtype=np.bool)\n",
        "    token_mask[np.array(token_mask_ids)] = True\n",
        "    self.token_mask = token_mask\n",
        "\n",
        "    self.start_token = index_from_string(tf.constant('[START]'))\n",
        "    self.end_token = index_from_string(tf.constant('[END]'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iBQzFZ9uWU79"
      },
      "outputs": [],
      "source": [
        "translator = Translator(\n",
        "    encoder=train_translator.encoder,\n",
        "    decoder=train_translator.decoder,\n",
        "    input_text_processor=input_text_processor,\n",
        "    output_text_processor=output_text_processor,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b59PN-UxqYrU"
      },
      "source": [
        "### Convert token IDs to text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-razg3Aso737"
      },
      "source": [
        "The first method to implement is `tokens_to_text` which converts from token IDs to human readable text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8IjwKTwtmdFf"
      },
      "outputs": [],
      "source": [
        "def tokens_to_text(self, result_tokens):\n",
        "  shape_checker = ShapeChecker()\n",
        "  shape_checker(result_tokens, ('batch', 't'))\n",
        "  result_text_tokens = self.output_token_string_from_index(result_tokens)\n",
        "  shape_checker(result_text_tokens, ('batch', 't'))\n",
        "\n",
        "  result_text = tf.strings.reduce_join(result_text_tokens,\n",
        "                                       axis=1, separator=' ')\n",
        "  shape_checker(result_text, ('batch'))\n",
        "\n",
        "  result_text = tf.strings.strip(result_text)\n",
        "  shape_checker(result_text, ('batch',))\n",
        "  return result_text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "912aV0K7r90w"
      },
      "outputs": [],
      "source": [
        "Translator.tokens_to_text = tokens_to_text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "krBuAapkqNs9"
      },
      "source": [
        "Input some random token IDs and see what it generates:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cWCMHdoS32QN"
      },
      "outputs": [],
      "source": [
        "example_output_tokens = tf.random.uniform(\n",
        "    shape=[5, 2], minval=0, dtype=tf.int64,\n",
        "    maxval=output_text_processor.vocabulary_size())\n",
        "translator.tokens_to_text(example_output_tokens).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AC9De_kAqtaE"
      },
      "source": [
        "### Sample from the decoder's predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q5tno-2ksJv6"
      },
      "source": [
        "This function takes the decoder's logit outputs and samples token IDs from that distribution:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8lfuj3GcdD6e"
      },
      "outputs": [],
      "source": [
        "def sample(self, logits, temperature):\n",
        "  shape_checker = ShapeChecker()\n",
        "  # 't' is usually 1 here.\n",
        "  shape_checker(logits, ('batch', 't', 'vocab'))\n",
        "  shape_checker(self.token_mask, ('vocab',))\n",
        "\n",
        "  token_mask = self.token_mask[tf.newaxis, tf.newaxis, :]\n",
        "  shape_checker(token_mask, ('batch', 't', 'vocab'), broadcast=True)\n",
        "\n",
        "  # Set the logits for all masked tokens to -inf, so they are never chosen.\n",
        "  logits = tf.where(self.token_mask, -np.inf, logits)\n",
        "\n",
        "  if temperature == 0.0:\n",
        "    new_tokens = tf.argmax(logits, axis=-1)\n",
        "  else: \n",
        "    logits = tf.squeeze(logits, axis=1)\n",
        "    new_tokens = tf.random.categorical(logits/temperature,\n",
        "                                        num_samples=1)\n",
        "  \n",
        "  shape_checker(new_tokens, ('batch', 't'))\n",
        "\n",
        "  return new_tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4DpDnBdBdL9_"
      },
      "outputs": [],
      "source": [
        "Translator.sample = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QwdHfGEfsmy5"
      },
      "source": [
        "Test run this function on some random inputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rwLT0nxXym80"
      },
      "outputs": [],
      "source": [
        "example_logits = tf.random.normal([5, 1, output_text_processor.vocabulary_size()])\n",
        "example_output_tokens = translator.sample(example_logits, temperature=1.0)\n",
        "example_output_tokens"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NEWIKFIJ2HWM"
      },
      "source": [
        "### Implement the translation loop\n",
        "\n",
        "Here is a complete implementation of the text to text translation loop.\n",
        "\n",
        "This implementation collects the results into python lists, before using `tf.concat` to join them into tensors.\n",
        "\n",
        "This implementation statically unrolls the graph out to `max_length` iterations.\n",
        "This is okay with eager execution in python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZmOvVrZmwAxg"
      },
      "outputs": [],
      "source": [
        "def translate_unrolled(self,\n",
        "                       input_text, *,\n",
        "                       max_length=50,\n",
        "                       return_attention=True,\n",
        "                       temperature=1.0):\n",
        "  batch_size = tf.shape(input_text)[0]\n",
        "  input_tokens = self.input_text_processor(input_text)\n",
        "  enc_output, enc_state = self.encoder(input_tokens)\n",
        "\n",
        "  dec_state = enc_state\n",
        "  new_tokens = tf.fill([batch_size, 1], self.start_token)\n",
        "\n",
        "  result_tokens = []\n",
        "  attention = []\n",
        "  done = tf.zeros([batch_size, 1], dtype=tf.bool)\n",
        "\n",
        "  for _ in range(max_length):\n",
        "    dec_input = DecoderInput(new_tokens=new_tokens,\n",
        "                             enc_output=enc_output,\n",
        "                             mask=(input_tokens!=0))\n",
        "    \n",
        "    dec_result, dec_state = self.decoder(dec_input, state=dec_state)\n",
        "\n",
        "    attention.append(dec_result.attention_weights)\n",
        "\n",
        "    new_tokens = self.sample(dec_result.logits, temperature)\n",
        "\n",
        "    # If a sequence produces an `end_token`, set it `done`\n",
        "    done = done | (new_tokens == self.end_token)\n",
        "    # Once a sequence is done it only produces 0-padding.\n",
        "    new_tokens = tf.where(done, tf.constant(0, dtype=tf.int64), new_tokens)\n",
        "\n",
        "    # Collect the generated tokens\n",
        "    result_tokens.append(new_tokens)\n",
        "\n",
        "    if tf.executing_eagerly() and tf.reduce_all(done):\n",
        "      break\n",
        "\n",
        "  # Convert the list of generates token ids to a list of strings.\n",
        "  result_tokens = tf.concat(result_tokens, axis=-1)\n",
        "  result_text = self.tokens_to_text(result_tokens)\n",
        "\n",
        "  if return_attention:\n",
        "    attention_stack = tf.concat(attention, axis=1)\n",
        "    return {'text': result_text, 'attention': attention_stack}\n",
        "  else:\n",
        "    return {'text': result_text}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JOmd8Y269MG3"
      },
      "outputs": [],
      "source": [
        "Translator.translate = translate_unrolled"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NxYXf3GNKKLS"
      },
      "source": [
        "Run it on a simple input:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hd2rgyHwVVrv"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "input_text = tf.constant([\n",
        "    'hace mucho frio aqui.', # \"It's really cold here.\"\n",
        "    'Esta es mi vida.', # \"This is my life.\"\"\n",
        "])\n",
        "\n",
        "result = translator.translate(\n",
        "    input_text = input_text)\n",
        "\n",
        "print(result['text'][0].numpy().decode())\n",
        "print(result['text'][1].numpy().decode())\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S-6cFyqeUPQm"
      },
      "source": [
        "If you want to export this model you'll need to wrap this method in a `tf.function`. This basic implementation has a few issues if you try to do that:\n",
        "\n",
        "1. The resulting graphs are very large and take a few seconds to build, save or load.\n",
        "2. You can't break from a statically unrolled loop, so it will always run `max_length` iterations, even if all the outputs are done. But even then it's marginally faster than eager execution.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_JhTZ5hOptO-"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=[tf.TensorSpec(dtype=tf.string, shape=[None])])\n",
        "def tf_translate(self, input_text):\n",
        "  return self.translate(input_text)\n",
        "\n",
        "Translator.tf_translate = tf_translate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fkccvHDvXCa8"
      },
      "source": [
        "Run the `tf.function` once to compile it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_NzrixLvVBjQ"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "result = translator.tf_translate(\n",
        "    input_text = input_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "USJdu00tVFbd"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "result = translator.tf_translate(\n",
        "    input_text = input_text)\n",
        "\n",
        "print(result['text'][0].numpy().decode())\n",
        "print(result['text'][1].numpy().decode())\n",
        "print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EbQpyYs13jF_"
      },
      "outputs": [],
      "source": [
        "#@title [Optional] Use a symbolic loop\n",
        "def translate_symbolic(self,\n",
        "                       input_text,\n",
        "                       *,\n",
        "                       max_length=50,\n",
        "                       return_attention=True,\n",
        "                       temperature=1.0):\n",
        "  shape_checker = ShapeChecker()\n",
        "  shape_checker(input_text, ('batch',))\n",
        "\n",
        "  batch_size = tf.shape(input_text)[0]\n",
        "\n",
        "  # Encode the input\n",
        "  input_tokens = self.input_text_processor(input_text)\n",
        "  shape_checker(input_tokens, ('batch', 's'))\n",
        "\n",
        "  enc_output, enc_state = self.encoder(input_tokens)\n",
        "  shape_checker(enc_output, ('batch', 's', 'enc_units'))\n",
        "  shape_checker(enc_state, ('batch', 'enc_units'))\n",
        "\n",
        "  # Initialize the decoder\n",
        "  dec_state = enc_state\n",
        "  new_tokens = tf.fill([batch_size, 1], self.start_token)\n",
        "  shape_checker(new_tokens, ('batch', 't1'))\n",
        "\n",
        "  # Initialize the accumulators\n",
        "  result_tokens = tf.TensorArray(tf.int64, size=1, dynamic_size=True)\n",
        "  attention = tf.TensorArray(tf.float32, size=1, dynamic_size=True)\n",
        "  done = tf.zeros([batch_size, 1], dtype=tf.bool)\n",
        "  shape_checker(done, ('batch', 't1'))\n",
        "\n",
        "  for t in tf.range(max_length):\n",
        "    dec_input = DecoderInput(\n",
        "        new_tokens=new_tokens, enc_output=enc_output, mask=(input_tokens != 0))\n",
        "\n",
        "    dec_result, dec_state = self.decoder(dec_input, state=dec_state)\n",
        "\n",
        "    shape_checker(dec_result.attention_weights, ('batch', 't1', 's'))\n",
        "    attention = attention.write(t, dec_result.attention_weights)\n",
        "\n",
        "    new_tokens = self.sample(dec_result.logits, temperature)\n",
        "    shape_checker(dec_result.logits, ('batch', 't1', 'vocab'))\n",
        "    shape_checker(new_tokens, ('batch', 't1'))\n",
        "\n",
        "    # If a sequence produces an `end_token`, set it `done`\n",
        "    done = done | (new_tokens == self.end_token)\n",
        "    # Once a sequence is done it only produces 0-padding.\n",
        "    new_tokens = tf.where(done, tf.constant(0, dtype=tf.int64), new_tokens)\n",
        "\n",
        "    # Collect the generated tokens\n",
        "    result_tokens = result_tokens.write(t, new_tokens)\n",
        "\n",
        "    if tf.reduce_all(done):\n",
        "      break\n",
        "\n",
        "  # Convert the list of generated token ids to a list of strings.\n",
        "  result_tokens = result_tokens.stack()\n",
        "  shape_checker(result_tokens, ('t', 'batch', 't0'))\n",
        "  result_tokens = tf.squeeze(result_tokens, -1)\n",
        "  result_tokens = tf.transpose(result_tokens, [1, 0])\n",
        "  shape_checker(result_tokens, ('batch', 't'))\n",
        "\n",
        "  result_text = self.tokens_to_text(result_tokens)\n",
        "  shape_checker(result_text, ('batch',))\n",
        "\n",
        "  if return_attention:\n",
        "    attention_stack = attention.stack()\n",
        "    shape_checker(attention_stack, ('t', 'batch', 't1', 's'))\n",
        "\n",
        "    attention_stack = tf.squeeze(attention_stack, 2)\n",
        "    shape_checker(attention_stack, ('t', 'batch', 's'))\n",
        "\n",
        "    attention_stack = tf.transpose(attention_stack, [1, 0, 2])\n",
        "    shape_checker(attention_stack, ('batch', 't', 's'))\n",
        "\n",
        "    return {'text': result_text, 'attention': attention_stack}\n",
        "  else:\n",
        "    return {'text': result_text}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ngywxv1WYO_O"
      },
      "outputs": [],
      "source": [
        "Translator.translate = translate_symbolic"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lItV7qjEGsYc"
      },
      "source": [
        "The initial implementation used python lists to collect the outputs. This uses `tf.range` as the loop iterator, allowing `tf.autograph` to convert the loop. The biggest change in this implementation is the use of `tf.TensorArray` instead of python `list` to accumulate tensors. `tf.TensorArray` is required to collect a variable number of tensors in graph mode. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AJ_NznOgZTxC"
      },
      "source": [
        "With eager execution this implementation performs on par with the original:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JRh66y-YYeBw"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "result = translator.translate(\n",
        "    input_text = input_text)\n",
        "\n",
        "print(result['text'][0].numpy().decode())\n",
        "print(result['text'][1].numpy().decode())\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l6B8W4_MZdX0"
      },
      "source": [
        "But when you wrap it in a `tf.function` you'll notice two differences."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WX6EF8KtYh20"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=[tf.TensorSpec(dtype=tf.string, shape=[None])])\n",
        "def tf_translate(self, input_text):\n",
        "  return self.translate(input_text)\n",
        "\n",
        "Translator.tf_translate = tf_translate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9S0kQ-bBZswZ"
      },
      "source": [
        "First: Graph creation is much faster (~10x), since it doesn't create `max_iterations` copies of the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Eq8d40RKYoJa"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "result = translator.tf_translate(\n",
        "    input_text = input_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ABEwtKIZ6eE"
      },
      "source": [
        "Second: The compiled function is much faster on small inputs (5x on this example), because it can break out of the loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5VdCLxPYrpz"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "result = translator.tf_translate(\n",
        "    input_text = input_text)\n",
        "\n",
        "print(result['text'][0].numpy().decode())\n",
        "print(result['text'][1].numpy().decode())\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eo5sf4jZaO2l"
      },
      "source": [
        "### Visualize the process"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FzZzC2cJacTv"
      },
      "source": [
        "The attention weights returned by the `translate` method show where the model was \"looking\" when it generated each output token.\n",
        "\n",
        "So the sum of the attention over the input should return all ones:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UEd2GljgqQ-0"
      },
      "outputs": [],
      "source": [
        "a = result['attention'][0]\n",
        "\n",
        "print(np.sum(a, axis=-1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k_HWQHcI2_h5"
      },
      "source": [
        "Here is the attention distribution for the first output step of the first example. Note how the attention is now much more focused than it was for the untrained model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M8BHdqQujALu"
      },
      "outputs": [],
      "source": [
        "_ = plt.bar(range(len(a[0, :])), a[0, :])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qB13OG472Z3V"
      },
      "source": [
        "Since there is some rough alignment between the input and output words, you expect the attention to be focused near the diagonal:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xyeXuEYHd0kQ"
      },
      "outputs": [],
      "source": [
        "plt.imshow(np.array(a), vmin=0.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mXECcNTn2mxN"
      },
      "source": [
        "Here is some code to make a better attention plot:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s5hQWlbN3jGF"
      },
      "outputs": [],
      "source": [
        "#@title Labeled attention plots\n",
        "def plot_attention(attention, sentence, predicted_sentence):\n",
        "  sentence = tf_lower_and_split_punct(sentence).numpy().decode().split()\n",
        "  predicted_sentence = predicted_sentence.numpy().decode().split() + ['[END]']\n",
        "  fig = plt.figure(figsize=(10, 10))\n",
        "  ax = fig.add_subplot(1, 1, 1)\n",
        "\n",
        "  attention = attention[:len(predicted_sentence), :len(sentence)]\n",
        "\n",
        "  ax.matshow(attention, cmap='viridis', vmin=0.0)\n",
        "\n",
        "  fontdict = {'fontsize': 14}\n",
        "\n",
        "  ax.set_xticklabels([''] + sentence, fontdict=fontdict, rotation=90)\n",
        "  ax.set_yticklabels([''] + predicted_sentence, fontdict=fontdict)\n",
        "\n",
        "  ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
        "  ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
        "\n",
        "  ax.set_xlabel('Input text')\n",
        "  ax.set_ylabel('Output text')\n",
        "  plt.suptitle('Attention weights')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rrGawQv2eiA4"
      },
      "outputs": [],
      "source": [
        "i=0\n",
        "plot_attention(result['attention'][i], input_text[i], result['text'][i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JHBdOf9duumm"
      },
      "source": [
        "Translate a few more sentences and plot them:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WrAM0FDomq3E"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "three_input_text = tf.constant([\n",
        "    # This is my life.\n",
        "    'Esta es mi vida.',\n",
        "    # Are they still home?\n",
        "    '¿Todavía están en casa?',\n",
        "    # Try to find out.'\n",
        "    'Tratar de descubrir.',\n",
        "])\n",
        "\n",
        "result = translator.tf_translate(three_input_text)\n",
        "\n",
        "for tr in result['text']:\n",
        "  print(tr.numpy().decode())\n",
        "\n",
        "print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-LjFp0AljOaZ"
      },
      "outputs": [],
      "source": [
        "result['text']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v7QwIMrG-id2"
      },
      "outputs": [],
      "source": [
        "i = 0\n",
        "plot_attention(result['attention'][i], three_input_text[i], result['text'][i])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zYVoVf8P-lr-"
      },
      "outputs": [],
      "source": [
        "i = 1\n",
        "plot_attention(result['attention'][i], three_input_text[i], result['text'][i])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9sFvlZBk-me4"
      },
      "outputs": [],
      "source": [
        "i = 2\n",
        "plot_attention(result['attention'][i], three_input_text[i], result['text'][i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rA3xI3NzrRJt"
      },
      "source": [
        "The short sentences often work well, but if the input is too long the model literally loses focus and stops providing reasonable predictions. There are two main reasons for this:\n",
        "\n",
        "1. The model was trained with teacher-forcing feeding the correct token at each step, regardless of the model's predictions. The model could be made more robust if it were sometimes fed its own predictions.\n",
        "2. The model only has access to its previous output through the RNN state. If the RNN state gets corrupted, there's no way for the model to recover. [Transformers](transformer.ipynb) solve this by using self-attention in the encoder and decoder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-FUHFLEvSMbG"
      },
      "outputs": [],
      "source": [
        "long_input_text = tf.constant([inp[-1]])\n",
        "\n",
        "import textwrap\n",
        "print('Expected output:\\n', '\\n'.join(textwrap.wrap(targ[-1])))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lDa_8NaN_RUy"
      },
      "outputs": [],
      "source": [
        "result = translator.tf_translate(long_input_text)\n",
        "\n",
        "i = 0\n",
        "plot_attention(result['attention'][i], long_input_text[i], result['text'][i])\n",
        "_ = plt.suptitle('This never works')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mMA9Pp71nzH9"
      },
      "source": [
        "## Export"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5rLMNOmsKoXe"
      },
      "source": [
        "Once you have a model you're satisfied with you might want to export it as a `tf.saved_model` for use outside of this python program that created it.\n",
        "\n",
        "Since the model is a subclass of `tf.Module` (through `keras.Model`), and all the functionality for export is compiled in a `tf.function` the model should export cleanly with `tf.saved_model.save`:  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NP2dNtEXJPEL"
      },
      "source": [
        "Now that the function has been traced it can be exported using `saved_model.save`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OyvxT5V0_X5B"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(translator, 'translator',\n",
        "                    signatures={'serving_default': translator.tf_translate})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-I0j3i3ekOba"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load('translator')\n",
        "result = reloaded.tf_translate(three_input_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GXZF__FZXJCm"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "result = reloaded.tf_translate(three_input_text)\n",
        "\n",
        "for tr in result['text']:\n",
        "  print(tr.numpy().decode())\n",
        "\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RTe5P5ioMJwN"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "* [Download a different dataset](http://www.manythings.org/anki/) to experiment with translations, for example, English to German, or English to French.\n",
        "* Experiment with training on a larger dataset, or using more epochs.\n",
        "* Try the [transformer tutorial](transformer.ipynb) which implements a similar translation task but uses a transformer layers instead of RNNs. This version also uses a `text.BertTokenizer` to implement wordpiece tokenization.\n",
        "* Have a look at the [tensorflow_addons.seq2seq](https://www.tensorflow.org/addons/tutorials/networks_seq2seq_nmt) for implementing this sort of sequence to sequence model. The `tfa.seq2seq` package includes higher level functionality like `seq2seq.BeamSearchDecoder`. "
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "nmt_with_attention.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
